"""
使用 LangChain 和 Qwen3 的企业数据处理示例

本模块展示了使用 LangChain 和 Qwen3-8B 模型（通过 Ollama）进行各种企业数据处理任务，
包括：
- 文本清理和标准化
- 数据分类和提取
- 文档摘要
- 关键信息提取
- 数据格式转换
- 情感分析
- 关键词提取
- 文本相似度比较
- 实体关系提取
- 文本翻译
"""

import json
import re
from typing import List, Dict, Any, Tuple

from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from langchain.text_splitter import CharacterTextSplitter
from langchain_community.llms import Ollama


class EnterpriseDataProcessor:
    def __init__(self):
        """初始化企业数据处理器，使用 Ollama 模型"""
        self.llm = Ollama(model="qwen3:8b", temperature=0.7)
        self.text_splitter = CharacterTextSplitter(
            separator="\n",
            chunk_size=2000,
            chunk_overlap=200
        )

    def clean_text(self, text: str, use_llm: bool = False) -> str:
        """
        清理和标准化文本
        
        参数:
            text: 输入文本
            use_llm: 是否使用 LLM 进行智能清理
            
        返回:
            清理后的文本
        """
        if not use_llm:
            # 基础清理：移除多余空白、标准化标点等
            text = re.sub(r'\s+', ' ', text)
            text = text.strip()
            return text

        # 使用 LLM 进行智能清理
        prompt = PromptTemplate(
            input_variables=["text"],
            template="请清理并标准化以下文本，保持其原始含义：\n{text}"
        )
        chain = LLMChain(llm=self.llm, prompt=prompt)
        return chain.run(text)

    def classify_document(self, text: str, categories: List[str]) -> str:
        """
        对文档进行分类
        
        参数:
            text: 输入文本
            categories: 可选的类别列表
            
        返回:
            最匹配的类别
        """
        prompt = PromptTemplate(
            input_variables=["text", "categories"],
            template="请将以下文本分类到给定的类别中：\n文本：{text}\n类别：{categories}\n最匹配的类别是："
        )
        chain = LLMChain(llm=self.llm, prompt=prompt)
        return chain.run({"text": text, "categories": categories})

    def extract_key_info(self, text: str) -> Dict[str, Any]:
        """
        从文本中提取关键信息
        
        参数:
            text: 输入文本
            
        返回:
            包含关键信息的字典
        """
        prompt = PromptTemplate(
            input_variables=["text"],
            template="请从以下文本中提取关键信息，以JSON格式返回：\n{text}"
        )
        chain = LLMChain(llm=self.llm, prompt=prompt)
        result = chain.run(text)
        try:
            return json.loads(result)
        except:
            return {"error": "无法解析JSON结果"}

    def summarize_document(self, text: str, max_length: int = 200) -> str:
        """
        生成文档摘要
        
        参数:
            text: 输入文本
            max_length: 摘要最大长度
            
        返回:
            文档摘要
        """
        # 对长文本进行分块处理
        if len(text) > 2000:
            chunks = self.text_splitter.split_text(text)
            summaries = []
            for chunk in chunks:
                prompt = PromptTemplate(
                    input_variables=["text", "max_length"],
                    template="请生成以下文本的摘要，不超过{max_length}字：\n{text}"
                )
                chain = LLMChain(llm=self.llm, prompt=prompt)
                summaries.append(chain.run({"text": chunk, "max_length": max_length}))
            return " ".join(summaries)

        prompt = PromptTemplate(
            input_variables=["text", "max_length"],
            template="请生成以下文本的摘要，不超过{max_length}字：\n{text}"
        )
        chain = LLMChain(llm=self.llm, prompt=prompt)
        return chain.run({"text": text, "max_length": max_length})

    def convert_format(self, text: str, target_format: str) -> str:
        """
        转换文本格式
        
        参数:
            text: 输入文本
            target_format: 目标格式（如：JSON、XML、CSV等）
            
        返回:
            转换后的文本
        """
        prompt = PromptTemplate(
            input_variables=["text", "target_format"],
            template="请将以下文本转换为{target_format}格式：\n{text}"
        )
        chain = LLMChain(llm=self.llm, prompt=prompt)
        return chain.run({"text": text, "target_format": target_format})

    def analyze_sentiment(self, text: str, aspect: str = None) -> Dict[str, Any]:
        """
        分析文本情感
        
        参数:
            text: 输入文本
            aspect: 可选的情感分析方面
            
        返回:
            情感分析结果
        """
        if aspect:
            prompt = PromptTemplate(
                input_variables=["text", "aspect"],
                template="请分析以下文本在{aspect}方面的情感倾向：\n{text}"
            )
            chain = LLMChain(llm=self.llm, prompt=prompt)
            return chain.run({"text": text, "aspect": aspect})

        prompt = PromptTemplate(
            input_variables=["text"],
            template="请分析以下文本的情感倾向：\n{text}"
        )
        chain = LLMChain(llm=self.llm, prompt=prompt)
        return chain.run(text)

    def extract_keywords(self, text: str, top_n: int = 5) -> List[Tuple[str, float]]:
        """
        提取文本关键词
        
        参数:
            text: 输入文本
            top_n: 返回的关键词数量
            
        返回:
            关键词及其重要性得分的列表
        """
        prompt = PromptTemplate(
            input_variables=["text", "top_n"],
            template="请从以下文本中提取{top_n}个关键词及其重要性得分：\n{text}"
        )
        chain = LLMChain(llm=self.llm, prompt=prompt)
        result = chain.run({"text": text, "top_n": top_n})
        try:
            return json.loads(result)
        except:
            return []

    def compare_texts(self, text1: str, text2: str) -> Dict[str, Any]:
        """
        比较两段文本的相似度
        
        参数:
            text1: 第一段文本
            text2: 第二段文本
            
        返回:
            相似度分析结果
        """
        prompt = PromptTemplate(
            input_variables=["text1", "text2"],
            template="请比较以下两段文本的相似度：\n文本1：{text1}\n文本2：{text2}"
        )
        chain = LLMChain(llm=self.llm, prompt=prompt)
        return chain.run({"text1": text1, "text2": text2})

    def extract_entity_relations(self, text: str) -> List[Dict[str, str]]:
        """
        提取文本中的实体关系
        
        参数:
            text: 输入文本
            
        返回:
            实体关系列表
        """
        prompt = PromptTemplate(
            input_variables=["text"],
            template="请从以下文本中提取实体及其关系：\n{text}"
        )
        chain = LLMChain(llm=self.llm, prompt=prompt)
        result = chain.run(text)
        try:
            return json.loads(result)
        except:
            return []

    def translate_text(self, text: str, target_lang: str) -> str:
        """
        翻译文本
        
        参数:
            text: 输入文本
            target_lang: 目标语言
            
        返回:
            翻译后的文本
        """
        prompt = PromptTemplate(
            input_variables=["text", "target_lang"],
            template="请将以下文本翻译成{target_lang}：\n{text}"
        )
        chain = LLMChain(llm=self.llm, prompt=prompt)
        return chain.run({"text": text, "target_lang": target_lang})

def main():
    """Example usage of EnterpriseDataProcessor"""
    processor = EnterpriseDataProcessor()
    
    # Example text
    sample_text = """
    北京智能科技有限公司（简称：智科）于2023年3月15日与上海创新技术股份有限公司（简称：创新科技）
    签署战略合作协议。根据协议，智科将投资2000万元人民币用于联合研发新一代人工智能算法。
    项目预计持续18个月，预期将为双方带来超过5000万元的经济效益。
    """
    
    # 示例1：文本清理
    cleaned_text = processor.clean_text(sample_text)
    print("清理后的文本:", cleaned_text)
    
    # 示例2：文档分类
    categories = ["商业合作", "技术研发", "投资公告", "人事变动"]
    category = processor.classify_document(sample_text, categories)
    print("文档类别:", category)
    
    # 示例3：信息提取
    info_types = ["公司名称", "日期", "投资金额", "预期收益"]
    extracted_info = processor.extract_key_info(sample_text)
    print("提取的信息:", json.dumps(extracted_info, ensure_ascii=False, indent=2))
    
    # 示例4：生成摘要
    summary = processor.summarize_document(sample_text, max_length=100)
    print("文档摘要:", summary)
    
    # 新功能示例
    print("\n=== 新功能示例 ===")
    
    # 示例5：情感分析
    sentiment = processor.analyze_sentiment(sample_text)
    print("\n情感分析结果:", json.dumps(sentiment, ensure_ascii=False, indent=2))
    
    # 示例6：关键词提取
    keywords = processor.extract_keywords(sample_text, top_n=5)
    print("\n关键词提取结果:", json.dumps(keywords, ensure_ascii=False, indent=2))
    
    # 示例7：实体关系提取
    relationships = processor.extract_entity_relations(sample_text)
    print("\n实体关系提取结果:", json.dumps(relationships, ensure_ascii=False, indent=2))
    
    # 示例8：文本翻译
    translated_text = processor.translate_text(sample_text, "英文")
    print("\n英文翻译结果:", translated_text)
    
    # 示例9：文本相似度比较
    another_text = """
    上海创新技术股份有限公司今日宣布与北京智能科技有限公司达成战略合作，
    双方将在人工智能领域展开深度合作，总投资达2000万元。
    """
    similarity = processor.compare_texts(sample_text, another_text)
    print("\n文本相似度分析:", json.dumps(similarity, ensure_ascii=False, indent=2))

if __name__ == "__main__":
    main() 